"""
WinAPI窗口捕获模块
使用Windows API确保能够捕获被遮挡的窗口内容
"""

import ctypes
from ctypes import wintypes
import numpy as np
from PIL import Image
import time
import win32gui
import win32ui
import win32con
import win32api
from ctypes import windll

# Windows API常量
GWL_EXSTYLE = -20
WS_EX_LAYERED = 0x00080000
WS_EX_TRANSPARENT = 0x00000020

class WinAPICapture:
    """使用WinAPI进行窗口捕获的类"""
    
    def __init__(self):
        self.user32 = windll.user32
        self.gdi32 = windll.gdi32
        self.kernel32 = windll.kernel32
        
        # 缓存设备上下文以提高性能
        self.cached_dcs = {}  # 缓存窗口设备上下文
        self.screen_dc_cache = None  # 缓存屏幕设备上下文
        self.last_screen_size = None  # 记录上次屏幕尺寸
        
        # 缓存窗口handle和相关资源
        self.window_cache = {}  # 缓存窗口的handle和DC资源
        self.last_cleanup_time = time.time()  # 上次清理时间
        
        # 记录窗口的最佳捕获方法
        self.capture_method_cache = {}  # {hwnd: 'bitblt' 或 'printwindow'}
        
    def _get_cached_window_resources(self, hwnd):
        """获取缓存的窗口资源，如果不存在则创建"""
        if hwnd not in self.window_cache:
            try:
                # 获取窗口矩形
                rect = self.get_window_rect(hwnd)
                if not rect or rect['width'] <= 0 or rect['height'] <= 0:
                    return None
                
                # 创建设备上下文
                hwndDC = win32gui.GetWindowDC(hwnd)
                mfcDC = win32ui.CreateDCFromHandle(hwndDC)
                saveDC = mfcDC.CreateCompatibleDC()
                
                # 创建位图
                saveBitMap = win32ui.CreateBitmap()
                saveBitMap.CreateCompatibleBitmap(mfcDC, rect['width'], rect['height'])
                saveDC.SelectObject(saveBitMap)
                
                # 缓存资源
                self.window_cache[hwnd] = {
                    'hwndDC': hwndDC,
                    'mfcDC': mfcDC,
                    'saveDC': saveDC,
                    'bitmap': saveBitMap,
                    'width': rect['width'],
                    'height': rect['height'],
                    'last_used': time.time()
                }
            except Exception as e:
                print(f"创建窗口缓存资源失败: {e}")
                return None
        
        # 更新使用时间
        self.window_cache[hwnd]['last_used'] = time.time()
        return self.window_cache[hwnd]
    
    def _cleanup_old_window_cache(self):
        """清理长时间未使用的窗口缓存"""
        current_time = time.time()
        if current_time - self.last_cleanup_time < 30:  # 30秒清理一次
            return
        
        self.last_cleanup_time = current_time
        keys_to_remove = []
        
        for hwnd, cache_info in self.window_cache.items():
            if current_time - cache_info['last_used'] > 60:  # 60秒未使用
                try:
                    # 清理资源
                    win32gui.DeleteObject(cache_info['bitmap'].GetHandle())
                    cache_info['saveDC'].DeleteDC()
                    cache_info['mfcDC'].DeleteDC()
                    win32gui.ReleaseDC(hwnd, cache_info['hwndDC'])
                except:
                    pass
                keys_to_remove.append(hwnd)
        
        for key in keys_to_remove:
            del self.window_cache[key]
    
    def get_window_handle(self, window_title):
        """根据窗口标题获取窗口句柄"""
        if not window_title:
            return None
            
        def enum_windows_callback(hwnd, windows):
            if win32gui.IsWindowVisible(hwnd):
                window_text = win32gui.GetWindowText(hwnd)
                if window_title.lower() in window_text.lower():
                    windows.append((hwnd, window_text))
            return True
        
        windows = []
        win32gui.EnumWindows(enum_windows_callback, windows)
        
        if windows:
            # 返回第一个匹配的窗口
            return windows[0][0]
        return None
    
    def get_window_rect(self, hwnd):
        """获取窗口矩形区域"""
        try:
            rect = win32gui.GetWindowRect(hwnd)
            return {
                'left': rect[0],
                'top': rect[1], 
                'right': rect[2],
                'bottom': rect[3],
                'width': rect[2] - rect[0],
                'height': rect[3] - rect[1]
            }
        except:
            return None
    
    def is_window_minimized(self, hwnd):
        """检查窗口是否最小化"""
        try:
            return win32gui.IsIconic(hwnd)
        except:
            return False
    
    def is_image_black(self, img, threshold=15):
        """检测图像是否全黑或接近全黑
        
        Args:
            img: PIL Image对象
            threshold: 亮度阈值（0-255），低于此值认为是黑色
            
        Returns:
            bool: True表示图像是黑色
        """
        if img is None:
            return True
        
        # 转换为numpy数组
        img_array = np.array(img)
        
        # 计算平均亮度
        mean_brightness = img_array.mean()
        
        return mean_brightness < threshold
    
    def capture_window_by_handle(self, hwnd):
        """智能捕获窗口：记忆每个窗口的最佳捕获方法"""
        try:
            t_start = time.time()
            
            # 检查窗口是否最小化
            if self.is_window_minimized(hwnd):
                print(f"窗口 {hwnd} 已最小化")
                return None
            
            # 获取窗口位置和大小
            left, top, right, bottom = win32gui.GetWindowRect(hwnd)
            width = right - left
            height = bottom - top
            
            if width <= 0 or height <= 0:
                print(f"窗口 {hwnd} 尺寸无效: {width}x{height}")
                return None
            
            # 获取窗口的设备上下文（DC）
            hwnd_dc = win32gui.GetWindowDC(hwnd)
            mfc_dc = win32ui.CreateDCFromHandle(hwnd_dc)
            save_dc = mfc_dc.CreateCompatibleDC()
            
            # 创建位图对象
            save_bitmap = win32ui.CreateBitmap()
            save_bitmap.CreateCompatibleBitmap(mfc_dc, width, height)
            save_dc.SelectObject(save_bitmap)
            
            # 检查是否已记录该窗口的最佳捕获方法
            cached_method = self.capture_method_cache.get(hwnd)
            
            if cached_method == 'printwindow':
                # 该窗口之前检测到需要用PrintWindow，直接使用
                print(f"[使用缓存方法: PrintWindow]")
                t0 = time.time()
                result = self.user32.PrintWindow(hwnd, save_dc.GetSafeHdc(), 3)
                t1 = time.time()
                
                if result != 0:
                    bmpinfo = save_bitmap.GetInfo()
                    bmpstr = save_bitmap.GetBitmapBits(True)
                    img = Image.frombuffer(
                        'RGB',
                        (bmpinfo['bmWidth'], bmpinfo['bmHeight']),
                        bmpstr, 'raw', 'BGRX', 0, 1
                    )
                    print(f"[PrintWindow: {(t1-t0)*1000:.2f}ms]")
                else:
                    print(f"PrintWindow失败")
                    img = None
                    
            else:
                # 第一次捕获该窗口，或之前用BitBlt成功过
                # 先尝试使用BitBlt（快速）
                t0 = time.time()
                save_dc.BitBlt((0, 0), (width, height), mfc_dc, (0, 0), win32con.SRCCOPY)
                t1 = time.time()
                
                # 将位图转换为 PIL Image
                bmpinfo = save_bitmap.GetInfo()
                bmpstr = save_bitmap.GetBitmapBits(True)
                
                img = Image.frombuffer(
                    'RGB',
                    (bmpinfo['bmWidth'], bmpinfo['bmHeight']),
                    bmpstr, 'raw', 'BGRX', 0, 1
                )
                
                # 检测是否是黑色画面
                if self.is_image_black(img):
                    print(f"检测到黑色画面，使用PrintWindow重新捕获...")
                    t2 = time.time()
                    
                    # 使用PrintWindow重新捕获
                    result = self.user32.PrintWindow(hwnd, save_dc.GetSafeHdc(), 3)
                    t3 = time.time()
                    
                    if result != 0:
                        # PrintWindow成功，重新获取位图数据
                        bmpstr = save_bitmap.GetBitmapBits(True)
                        img = Image.frombuffer(
                            'RGB',
                            (bmpinfo['bmWidth'], bmpinfo['bmHeight']),
                            bmpstr, 'raw', 'BGRX', 0, 1
                        )
                        # 记录该窗口需要用PrintWindow
                        self.capture_method_cache[hwnd] = 'printwindow'
                        print(f"[已记忆: 窗口 {hwnd} 使用PrintWindow]")
                        print(f"[BitBlt: {(t1-t0)*1000:.2f}ms] [PrintWindow: {(t3-t2)*1000:.2f}ms]")
                    else:
                        print(f"PrintWindow也失败了，返回黑色画面")
                else:
                    # BitBlt成功，记录该窗口可以用BitBlt
                    if cached_method is None:
                        self.capture_method_cache[hwnd] = 'bitblt'
                        print(f"[已记忆: 窗口 {hwnd} 使用BitBlt]")
                    print(f"[BitBlt: {(t1-t0)*1000:.2f}ms]")
            
            # 清理资源
            win32gui.DeleteObject(save_bitmap.GetHandle())
            save_dc.DeleteDC()
            mfc_dc.DeleteDC()
            win32gui.ReleaseDC(hwnd, hwnd_dc)
            
            t_end = time.time()
            # print(f"[总耗时] {(t_end-t_start)*1000:.2f}ms")
            
            return img
            
        except Exception as e:
            print(f"WinAPI捕获窗口失败: {e}")
            return None
    
    def capture_window_by_title(self, window_title):
        """根据窗口标题捕获窗口内容"""
        hwnd = self.get_window_handle(window_title)
        if hwnd:
            return self.capture_window_by_handle(hwnd)
        return None
    
    def capture_screen(self):
        """捕获整个屏幕（优化版本）"""
        try:
            # 获取屏幕尺寸
            width = win32api.GetSystemMetrics(win32con.SM_CXSCREEN)
            height = win32api.GetSystemMetrics(win32con.SM_CYSCREEN)
            
            # 检查屏幕尺寸是否变化
            current_screen_size = (width, height)
            if (self.last_screen_size != current_screen_size or 
                self.screen_dc_cache is None):
                
                # 清理旧的缓存
                if self.screen_dc_cache:
                    try:
                        win32gui.DeleteObject(self.screen_dc_cache['bitmap'].GetHandle())
                        self.screen_dc_cache['saveDC'].DeleteDC()
                        self.screen_dc_cache['mfcDC'].DeleteDC()
                        win32gui.ReleaseDC(self.screen_dc_cache['hwnd'], self.screen_dc_cache['hwndDC'])
                    except:
                        pass
                
                # 创建新的设备上下文
                hwnd = win32gui.GetDesktopWindow()
                hwndDC = win32gui.GetWindowDC(hwnd)
                mfcDC = win32ui.CreateDCFromHandle(hwndDC)
                saveDC = mfcDC.CreateCompatibleDC()
                
                # 创建位图
                saveBitMap = win32ui.CreateBitmap()
                saveBitMap.CreateCompatibleBitmap(mfcDC, width, height)
                saveDC.SelectObject(saveBitMap)
                
                # 缓存设备上下文
                self.screen_dc_cache = {
                    'hwnd': hwnd,
                    'hwndDC': hwndDC,
                    'mfcDC': mfcDC,
                    'saveDC': saveDC,
                    'bitmap': saveBitMap,
                    'width': width,
                    'height': height
                }
                self.last_screen_size = current_screen_size
            
            # 使用缓存的设备上下文
            cache = self.screen_dc_cache
            
            # 复制屏幕内容
            cache['saveDC'].BitBlt((0, 0), (cache['width'], cache['height']), 
                                 cache['mfcDC'], (0, 0), win32con.SRCCOPY)
            
            # 获取位图数据
            bmpinfo = cache['bitmap'].GetInfo()
            bmpstr = cache['bitmap'].GetBitmapBits(True)
            
            # 转换为PIL图像
            img = Image.frombuffer(
                'RGB',
                (bmpinfo['bmWidth'], bmpinfo['bmHeight']),
                bmpstr, 'raw', 'BGRX', 0, 1
            )
            
            return img
            
        except Exception as e:
            print(f"WinAPI捕获屏幕失败: {e}")
            return None
    
    def clear_window_cache(self):
        """清理所有窗口缓存资源"""
        for hwnd, cache_info in self.window_cache.items():
            try:
                win32gui.DeleteObject(cache_info['bitmap'].GetHandle())
                cache_info['saveDC'].DeleteDC()
                cache_info['mfcDC'].DeleteDC()
                win32gui.ReleaseDC(hwnd, cache_info['hwndDC'])
            except:
                pass
        self.window_cache.clear()
        print("窗口缓存已清理")
    
    def get_cache_stats(self):
        """获取缓存统计信息"""
        bitblt_count = sum(1 for method in self.capture_method_cache.values() if method == 'bitblt')
        printwindow_count = sum(1 for method in self.capture_method_cache.values() if method == 'printwindow')
        
        return {
            'cached_windows': len(self.window_cache),
            'screen_cached': self.screen_dc_cache is not None,
            'method_cache': {
                'total': len(self.capture_method_cache),
                'bitblt': bitblt_count,
                'printwindow': printwindow_count
            }
        }
    
    def clear_method_cache(self, hwnd=None):
        """清除捕获方法缓存
        
        Args:
            hwnd: 窗口句柄，None表示清除所有
        """
        if hwnd is None:
            self.capture_method_cache.clear()
            print("已清除所有窗口的捕获方法缓存")
        elif hwnd in self.capture_method_cache:
            del self.capture_method_cache[hwnd]
            print(f"已清除窗口 {hwnd} 的捕获方法缓存")
    
    def get_all_windows(self):
        """获取所有可见窗口的列表"""
        windows = []
        
        def enum_windows_callback(hwnd, windows):
            if win32gui.IsWindowVisible(hwnd):
                window_text = win32gui.GetWindowText(hwnd)
                if window_text:  # 只包含有标题的窗口
                    try:
                        rect = self.get_window_rect(hwnd)
                        if rect and rect['width'] > 0 and rect['height'] > 0:
                            windows.append({
                                'hwnd': hwnd,
                                'title': window_text,
                                'rect': rect
                            })
                    except:
                        pass
            return True
        
        win32gui.EnumWindows(enum_windows_callback, windows)
        return windows

# 全局实例
winapi_capture = WinAPICapture()

def capture_window_winapi(window_title=None, is_fullscreen=False):
    """
    使用WinAPI捕获窗口或屏幕
    
    Args:
        window_title: 窗口标题，None表示全屏
        is_fullscreen: 是否全屏捕获
        
    Returns:
        PIL.Image对象或None
    """
    if is_fullscreen or not window_title:
        return winapi_capture.capture_screen()
    else:
        return winapi_capture.capture_window_by_title(window_title)

def get_all_windows_winapi():
    """获取所有可见窗口列表"""
    return winapi_capture.get_all_windows()

def test_winapi_capture():
    """测试WinAPI捕获功能"""
    print("测试WinAPI窗口捕获...")
    
    # 获取所有窗口
    windows = get_all_windows_winapi()
    print(f"找到 {len(windows)} 个窗口:")
    
    for i, window in enumerate(windows[:10]):  # 只显示前10个
        print(f"{i+1}. {window['title']} - {window['rect']['width']}x{window['rect']['height']}")
    
    # 测试捕获第一个窗口
    if windows:
        test_window = windows[0]
        print(f"\n测试捕获窗口: {test_window['title']}")
        
        # 第一次捕获（创建缓存）
        import time
        start_time = time.time()
        img = winapi_capture.capture_window_by_handle(test_window['hwnd'])
        first_capture_time = time.time() - start_time
        
        if img:
            print(f"捕获成功: {img.size}")
            img.save("test_winapi_capture.png")
            print("图像已保存为 test_winapi_capture.png")
            
            # 第二次捕获（使用缓存）
            start_time = time.time()
            img2 = winapi_capture.capture_window_by_handle(test_window['hwnd'])
            second_capture_time = time.time() - start_time
            
            print(f"第一次捕获时间: {first_capture_time:.3f}秒")
            print(f"第二次捕获时间: {second_capture_time:.3f}秒")
            print(f"性能提升: {first_capture_time/second_capture_time:.1f}x")
            
            # 显示缓存统计
            stats = winapi_capture.get_cache_stats()
            print(f"缓存统计: {stats}")
        else:
            print("捕获失败")

if __name__ == "__main__":
    test_winapi_capture()
